import os
import numpy as np
from utils.args_config import get_args

args = get_args()
print(os.getcwd())
latmin, latmax, lonmin, lonmax = 32.05, 43.2, 107.7, 117.0
width, height = 224, 184

rate = 20
C_to_K = np.float32(273.15)

National_station_ids = open('../data_verification/shanxi_national_stations.txt').readlines()
National_station_ids = [idx.strip() for idx in National_station_ids]


def cldas_points_score(Tem_cls):
    data_list = os.listdir(cldas_dir + Tem_cls)
    data_list.sort(key=lambda aa: os.path.basename(aa.strip()))

    for i, file in enumerate(data_list):
        # for file in ["202008170800.npy"]:
        score = []
        # if os.path.basename(file)[8:12] != "0800":
        #     continue

        cldas_file = os.path.join(cldas_dir, Tem_cls, file.strip())
        cldas_data = np.load(cldas_file).copy()

        cldas_time = os.path.basename(file).strip()[:12]
        station_time = cldas_time
        station_file = f'{args.stDataPath}/micaps{station_time[:4]}/micaps{station_time[:6]}' \
                       f'/SURFACE/TMP_{Tem_cls[1:4].upper()}_24H_NATIONAL/{station_time}00.000'
        if not os.path.isfile(station_file):
            os.system(f"cp {cldas_file} {revise_dir}/{Tem_cls}/")
            continue
        try:
            station_data_list = open(station_file, encoding='GBK').readlines()[5:]
        except UnicodeDecodeError:
            os.system(f"cp {cldas_file} {revise_dir}/{Tem_cls}/")
            continue
        station_cnt = 0
        for line in station_data_list:
            data = line.strip().split()
            assert len(data) == 5, "WRONG STATION DATA"
            if data[4] == '':
                continue
            station_id = data[0]
            # if station_id not in National_station_ids: continue
            Latitude = float(data[2])
            Longitude = float(data[1])
            TEM = float(data[4])
            # print(data_time, TEM_Max_24h, TEM_Min_24h, Latitude, Longitude)

            x_axis = int((latmax - Latitude) * rate + 0.5)
            y_axis = int((Longitude - lonmin) * rate + 0.5)

            if x_axis >= width or y_axis >= height or x_axis < 0 or y_axis < 0:
                continue
            station_cnt += 1
            TEM_24h_K = TEM + args.ck
            # 只改变站点位置
            cldas_data[x_axis, y_axis] = np.around(TEM_24h_K, 2)

            # 改变站点及附近一个点位置
            # nearCnt = 1
            # fst_data = cldas_data[x_axis, y_axis]
            # diff = TEM_24h_K - fst_data
            # re_w0 = max(0, x_axis - nearCnt)
            # re_w1 = min(width, x_axis + nearCnt + 1)
            # re_h0 = max(0, y_axis - nearCnt)
            # re_h1 = min(height, y_axis + nearCnt + 1)
            #
            # cldas_data[re_w0:re_w1, re_h0:re_h1] += diff

        re_dir = f'{revise_dir}/{Tem_cls}'
        if not os.path.isdir(re_dir):
            os.mkdir(re_dir)
        re_path = f'{revise_dir}/{Tem_cls}/{os.path.basename(file)}'

        np.save(re_path, cldas_data)
        print(i + 1, station_cnt, re_path)


if __name__ == '__main__':
    cldas_dir = "/home/gym/datasets/shanxi/temperature/target/"
    revise_dir = '/home/gym/datasets/shanxi/temperature-perday/re_cldas_553'
    if not os.path.isdir(revise_dir):
        os.mkdir(revise_dir)

    for tcls in ["TMIN"]:
        cldas_points_score(Tem_cls=tcls)
    # cldas_points_score(Tem_cls="TMIN")
